import datetime
import os
import pprint
import time
import threading
import torch as th
from types import SimpleNamespace as SN
from utils.logging import Logger
from utils.timehelper import time_left, time_str
from os.path import dirname, abspath

from learners import REGISTRY as le_REGISTRY
from runners import REGISTRY as r_REGISTRY
from controllers import REGISTRY as mac_REGISTRY
from components.episode_buffer import ReplayBuffer
from components.transforms import OneHot

import json
from copy import deepcopy

def run(_run, _config, _log):

    # check args sanity
    _config = args_sanity_check(_config, _log)

    args = SN(**_config)
    args.device = "cuda" if args.use_cuda else "cpu"

    # setup loggers
    logger = Logger(_log)

    _log.info("Experiment Parameters:")
    experiment_params = pprint.pformat(_config,
                                       indent=4,
                                       width=1)
    _log.info("\n\n" + experiment_params + "\n")

    # configure tensorboard logger
    results_save_dir = args.results_save_dir

    if args.use_tensorboard and not args.evaluate:
        # only log tensorboard when in training mode
        tb_exp_direc = os.path.join(results_save_dir, 'tb_logs')
        logger.setup_tb(tb_exp_direc)

        config_str = json.dumps(vars(args), indent=4)
        with open(os.path.join(results_save_dir, "config.json"), "w") as f:
            f.write(config_str)

    # sacred is on by default
    logger.setup_sacred(_run)

    # Run and train
    run_sequential(args=args, logger=logger)

    # Clean up after finishing
    print("Exiting Main")

    print("Stopping all threads")
    for t in threading.enumerate():
        if t.name != "MainThread":
            print("Thread {} is alive! Is daemon: {}".format(t.name, t.daemon))
            t.join(timeout=1)
            print("Thread joined")

    print("Exiting script")

    # Making sure framework really exits
    os._exit(os.EX_OK)

def evaluate_sequential(args, runner, mac_1, mac_2):

    ret_lst = []
    ret_additional_lst = []
    food_lst = []
    traj_lst = []
    for i in range(args.test_nepisode):
        ret, ret_additional, food, traj = runner.run_evaluate(mac_1, mac_2, i)
        ret_lst.append(ret)
        ret_additional_lst.append(ret_additional)
        food_lst.append(food)
        traj_lst.append(traj)
    from collections import Counter
    # print(Counter(food_lst))
    res = SN()
    res.ego = args.checkpoint_path_ego
    res.tm = args.checkpoint_path_tm
    res.food_lst = food_lst
    res.res = Counter(food_lst)
    import numpy as np
    res.ret = np.mean(ret_lst)
    res.ret_additional = np.array(ret_additional_lst).mean(axis=0).tolist()
    # res.reward_idx = int(args.env_args['key'].split('-')[-2].split('r')[1])
    res = json.dumps(vars(res), indent=4)
    with open(os.path.join(args.results_save_dir, "res.json"), "w") as f:
        f.write(res)

    if args.save_replay and args.save_traj:
        states = [traj['state'] for traj in traj_lst] # [1, seq_len, state_shape]
        states = th.cat(states, dim=0).cpu().numpy() # [n_traj, seq_len, state_shape]
        # print(states)
        np.save(os.path.join(args.results_save_dir, "states.npy"), states)

    if args.save_replay and not args.save_traj:
        runner.save_replay()

    runner.close_env()

def evaluate_sequential_multi_head(args, runner, mac_1, mac_2):
    for head_id in mac_1.agent.head_dict.keys():
        ret_lst = []
        ret_additional_lst = []
        food_lst = []
        traj_lst = []
        for i in range(args.test_nepisode):
            ret, ret_additional, food, traj = runner.run_evaluate(mac_1, mac_2, i, head_id)
            ret_lst.append(ret)
            ret_additional_lst.append(ret_additional)
            food_lst.append(food)
            traj_lst.append(traj)
        from collections import Counter
        # print(Counter(food_lst))
        res = SN()
        res.ego = args.checkpoint_path_ego
        res.tm = args.checkpoint_path_tm
        res.food_lst = food_lst
        res.res = Counter(food_lst)
        import numpy as np
        res.ret = np.mean(ret_lst)
        res.ret_additional = np.array(ret_additional_lst).mean(axis=0).tolist()
        # res.reward_idx = int(args.env_args['key'].split('-')[-2].split('r')[1])
        res = json.dumps(vars(res), indent=4)
        with open(os.path.join(args.results_save_dir, f"res_{head_id}.json"), "w") as f:
            f.write(res)

        if args.save_replay and args.save_traj:
            states = [traj['state'][:, :-1] for traj in traj_lst] # [1, seq_len, state_shape]
            states = th.cat(states, dim=0).cpu().numpy() # [n_traj, seq_len, state_shape]
            np.save(os.path.join(args.results_save_dir, f"states_{head_id}.npy"), states)

    runner.close_env()

def run_sequential(args, logger):

    # Init runner so we can get env info
    runner = r_REGISTRY[args.runner](args=args, logger=logger)

    # Set up schemes and groups here
    env_info = runner.get_env_info()
    args.n_agents = env_info["n_agents"]
    args.n_actions = env_info["n_actions"]
    args.state_shape = env_info["state_shape"]

    # Default/Base scheme
    scheme = {
        "state": {"vshape": env_info["state_shape"]},
        "obs": {"vshape": env_info["obs_shape"], "group": "agents"},
        "actions": {"vshape": (1,), "group": "agents", "dtype": th.long},
        "avail_actions": {"vshape": (env_info["n_actions"],), "group": "agents", "dtype": th.int},
        "reward": {"vshape": (1,)},
        "terminated": {"vshape": (1,), "dtype": th.uint8},
    }
    groups = {
        "agents": args.n_agents
    }
    preprocess = {
        "actions": ("actions_onehot", [OneHot(out_dim=args.n_actions)])
    }

    runner.setup(scheme=scheme, groups=groups, preprocess=preprocess)
    empty_buffer = ReplayBuffer(scheme, groups, args.buffer_size, env_info["episode_limit"] + 1,
                          preprocess=preprocess,
                          device="cpu" if args.buffer_cpu_only else args.device)

    assert args.checkpoint_path_ego != "" and args.checkpoint_path_tm != ""

    args.multi_head = 'multihead' in args.checkpoint_path_ego or 'multi_head' in args.checkpoint_path_ego

    if args.multi_head:
        mac_ego = mac_REGISTRY['multi_head_mac'](empty_buffer.scheme, groups, args)
        learner_ego = le_REGISTRY['multi_head_learner'](mac_ego, empty_buffer.scheme, logger, args)
    else:
        mac_ego = mac_REGISTRY[args.mac](empty_buffer.scheme, groups, args)
        learner_ego = le_REGISTRY[args.learner](mac_ego, empty_buffer.scheme, logger, args)
    if args.use_cuda:
        learner_ego.cuda()

    mac_tm = mac_REGISTRY[args.mac](empty_buffer.scheme, groups, args)
    learner_tm = le_REGISTRY[args.learner](mac_tm, empty_buffer.scheme, logger, args)
    if args.use_cuda:
        learner_tm.cuda()
        
    def get_timestep(path):
        timesteps = []
        timestep_to_load = 0
        if not os.path.isdir(path):
            logger.console_logger.info("Checkpoint directiory {} doesn't exist".format(args.checkpoint_path))
            return
        # Go through all files in args.checkpoint_path
        for name in os.listdir(path):
            full_name = os.path.join(path, name)
            # Check if they are dirs the names of which are numbers
            if os.path.isdir(full_name) and name.isdigit():
                timesteps.append(int(name))
        if args.load_step == 0:
            # choose the max timestep
            timestep_to_load = max(timesteps)
        else:
            # choose the timestep closest to load_step
            timestep_to_load = min(timesteps, key=lambda x: abs(x - args.load_step))
        return timestep_to_load

    timestep_to_load_ego = get_timestep(args.checkpoint_path_ego)
    timestep_to_load_tm  = get_timestep(args.checkpoint_path_tm)
    model_path_ego = os.path.join(args.checkpoint_path_ego, str(timestep_to_load_ego))
    model_path_tm = os.path.join(args.checkpoint_path_tm, str(timestep_to_load_tm))
    logger.console_logger.info("Loading ego model from {}".format(model_path_ego))
    logger.console_logger.info("Loading tm model from {}".format(model_path_tm))
    learner_ego.load_models(model_path_ego)
    learner_tm.load_models(model_path_tm)

    if args.multi_head:
        evaluate_sequential_multi_head(args, runner, mac_ego, mac_tm)
    else:
        evaluate_sequential(args, runner, mac_ego, mac_tm)
    return

def args_sanity_check(config, _log):

    # set CUDA flags
    # config["use_cuda"] = True # Use cuda whenever possible!
    if config["use_cuda"] and not th.cuda.is_available():
        config["use_cuda"] = False
        _log.warning("CUDA flag use_cuda was switched OFF automatically because no CUDA devices are available!")

    if config["test_nepisode"] < config["batch_size_run"]:
        config["test_nepisode"] = config["batch_size_run"]
    else:
        config["test_nepisode"] = (config["test_nepisode"]//config["batch_size_run"]) * config["batch_size_run"]

    return config